package edu.caltech.android.gaussianMixture;

import java.util.Vector;

import Jama.Matrix;

/**
 * @author jkrause
 */
public final class GaussianMixture {

	// private static final long serialVersionUID = -6889636018440660655L;

	/**
	 * Stores the weight assigned to each Gaussian. All weights should be
	 * positive and the sum of the weights should be 1.
	 */
	private final Vector<Double> weights;

	/**
	 * Stores the means of each Gaussian. Each mean should be a column vector of
	 * the same size, 'dim'.
	 */
	private final Vector<Matrix> means;

	/**
	 * Stores the covariances of each Gaussian. Each covariance should be
	 * positive definite and of size 'dim' by 'dim'.
	 */
	private final Vector<Matrix> covariances;

	// inverse of the covariance matrix of each Gaussian.
	private final Vector<Matrix> covInvs;

	// determinant of the covariance matrix of each Gaussian.
	private final Vector<Double> covDets;

	/**
	 * Number of Gaussians that this mixture is comprised of. Should be a
	 * positive integer.
	 */
	private final int numGaussians;

	// Dimension of each point.
	private final int dim;

	/**
	 * Constructor.
	 * 
	 * @param inWeights
	 *            Vector of weights for this mixture, specifying the portion of
	 *            the pdf made up of each Gaussian.
	 * @param inMeans
	 *            Vector of means for this mixture, giving the mean column
	 *            vector for each of the Gaussians.
	 * @param inCovs
	 *            Vector of covariance matrices for this mixture, giving the
	 *            covariance of each of the Gaussians.
	 */
	public GaussianMixture(final Vector<Double> inWeights,
			final Vector<Matrix> inMeans, final Vector<Matrix> inCovs) {
		// Must have the same number of mixtures for each parameter
		if (inWeights.size() != inMeans.size()
				|| inWeights.size() != inCovs.size()) {
			throw (new RuntimeException(
					"Mismatched dimensions: inWeights.size()="
							+ inWeights.size() + ", inMeans.size()="
							+ inMeans.size() + ", inCovs.size()="
							+ inCovs.size()));
		}

		numGaussians = inWeights.size();

		// Should have same dimension for everything
		dim = inMeans.firstElement().getRowDimension();

		// Check all mean vectors
		// Iterator<Matrix> iter = inMeans.iterator();
		// while (iter.hasNext()) {
		// final Matrix m = iter.next();
		// assert (m.getColumnDimension() == 1); // Should be a vector
		// assert (m.getRowDimension() == this.dim);
		// }
		for (Matrix m : inMeans) {
			if (m.getColumnDimension() != 1) {
				throw (new RuntimeException("mean is not a vector"));
			}
			if (m.getRowDimension() != dim) {
				throw (new RuntimeException("mean has wrong dimensionality"));
			}
		}

		// Check all covariance matrices
		// iter = inCovs.iterator();
		// while (iter.hasNext()) {
		// final Matrix m = iter.next();
		// assert (m.getColumnDimension() == this.dim); // Square matrix
		// assert (m.getRowDimension() == this.dim);
		// }
		for (Matrix m : inCovs) {
			if (m.getColumnDimension() != dim) {
				throw (new RuntimeException("covs is not a square matrix."));
			}
			if (m.getRowDimension() != dim) {
				throw (new RuntimeException("mean has number of rows."));
			}
		}

		// Make sure each weight is non-negative and that they sum to 1
		// final Iterator<Double> iterD = inWeights.iterator();
		// double weightsSum = 0;
		// while (iterD.hasNext()) {
		// final double weight = iterD.next().doubleValue();
		// assert (weight >= 0); // Non-negative
		// weightsSum += weight;
		// }

		double weightsSum = 0;
		for (double weight : inWeights) {
			if (weight < 0) {
				throw (new RuntimeException("Negative weight."));
			}
			weightsSum += weight;
		}

		// TODO: shouldn't this have a floating-point tolerance?
		// assert (weightsSum == 1); // Sums to 1

		double epsilon = 0.001;
		if (Math.abs(weightsSum - 1) > epsilon) {
			throw (new RuntimeException("weights do not sum to one."));
		}

		// Everything checks out
		weights = inWeights;
		means = inMeans;
		covariances = inCovs;

		covDets = new Vector<Double>();
		covInvs = new Vector<Matrix>();
		// Iterator<Matrix> iterC = this.covariances.iterator();
		// while (iterC.hasNext()) {
		// Matrix c = iterC.next();
		// covDets.add(new Double(c.det()));
		// covInvs.add(c.inverse());
		// }

		for (Matrix c : covariances) {
			covDets.add(new Double(c.det()));
			covInvs.add(c.inverse());
		}
	}

	/**
	 * Calculates the pdf of this Gaussian mixture at the given point 'x'.
	 * 
	 * @param x
	 *            A column vector giving the point at which to get the pdf.
	 * @return The value of the probability density function at that point.
	 */
	public final double pdf(final Matrix x) {
		// Make sure x is a vector of the appropriate length
		assert (x.getColumnDimension() == 1);
		assert (x.getRowDimension() == this.dim);

		// Calculate covariance determinants and inverses, if necessary.
		// Waiting until now to calculate them means that anything using this
		// as a storage class won't need to calculate them.
		// if (this.covDets == null) {
		// this.updateCovariances();
		// }

		// Sum up probabilities from each Gaussian
		double probSum = 0.;
		for (int i = 0; i < this.numGaussians; i++) {
			probSum += this.weights.get(i)
					* MathLib.multivariatePDF(x, this.means.get(i),
							this.covInvs.get(i), this.covDets.get(i)
									.doubleValue());
		}
		return probSum;
	}

	/**
	 * Calculates roughly the log pdf of this Gaussian mixture at the given
	 * point 'x'.
	 * 
	 * @param x
	 *            A column vector giving the point at which to get the pdf.
	 * @return The value of the probability density function at that point.
	 */
	public final double logPDF(final Matrix x) {
		// Make sure x is a vector of the appropriate length
		assert (x.getColumnDimension() == 1);
		assert (x.getRowDimension() == this.dim);

		// Calculate covariance determinants and inverses, if necessary.
		// Waiting until now to calculate them means that anything using this
		// as a storage class won't need to calculate them.
		// if (this.covDets == null) {
		// this.updateCovariances();
		// }

		// Sum up probabilities from each Gaussian
		double maxLog = Double.NEGATIVE_INFINITY;
		for (int i = 0; i < this.numGaussians; i++) {
			double log = Math.log(this.weights.get(i))
					+ MathLib.multivariateLogPDF(x, this.means.get(i),
							this.covInvs.get(i), this.covDets.get(i)
									.doubleValue());
			maxLog = Math.max(log, maxLog);
		}
		return maxLog;
	}

	/**
	 * @return The number of Gaussians this Gaussian Mixture is comprised of.
	 */
	public final int getNumGaussians() {
		return numGaussians;
	}

	/**
	 * Gets a view of the weight of each Gaussian in this GaussianMixture.
	 * 
	 * @return A view of the weight of each Gaussian in this GaussianMixture.
	 * @deprecated
	 */
	public final Vector<Double> getWeights() {
		return weights;
	}

	/**
	 * Gets a view of the mean of each Gaussian in this GaussianMixture.
	 * 
	 * @return A view of the mean of each Gaussian in this GaussianMixture.
	 * @deprecated
	 */
	public final Vector<Matrix> getMeans() {
		return means;
	}

	/**
	 * @return The dimension of this multivariate Gaussian.
	 */
	public final int getDimension() {
		return this.dim;
	}

	/**
	 * Gets a view of the covariance matrix of each Gaussian in this
	 * GaussianMixture. NOTE: IF THE COVARIANCES ARE CHANGED EXTERNALLY, THEN
	 * pdf WILL NO LONGER BEHAVE PROPERLY UNTIL A CALL TO updateCovariances IS
	 * MADE.
	 * 
	 * @return A view of the covariance matrix of each Gaussian in this
	 *         GaussianMixture.
	 * @deprecated
	 */
	public final Vector<Matrix> getCovariances() {
		return covariances;
	}

	/**
	 * Updates stored determinant and inverse of each covariance matrix.
	 */
	// public final void updateCovariances() {
	// this.covDets = new Vector<Double>();
	// this.covInvs = new Vector<Matrix>();
	// Iterator<Matrix> iter = this.covariances.iterator();
	// while (iter.hasNext()) {
	// Matrix c = iter.next();
	// covDets.add(new Double(c.det()));
	// covInvs.add(c.inverse());
	// }
	// }

}
